import random
import numpy as np
from envs.utils import silence_stderr
from envs.smacv1.StarCraft2_Env import StarCraft2Env as SMACEnv


class Config:

    add_local_obs = False
    add_move_state = False
    add_visible_state = False
    add_distance_state = False
    add_xy_state = False
    add_enemy_action_state = False
    add_agent_id = False
    use_state_agent = False
    use_mustalive = True
    add_center_xy = False
    use_stacked_frames = False
    stacked_frames = 1
    use_obs_instead_of_state = False
    n_eval_rollout_threads = 1
    env_name = "StarCraft2"

    def __init__(self, map_name, seed):
        self.map_name = map_name
        self.seed = seed


class SMACWrapper:
    
    def __init__(self, map_name, seed=0):
        random.seed(seed)
        np.random.seed(seed)
        # self.config = Config(map_name, seed)
        with silence_stderr():
            # self.env = SMACEnv(self.config)
            # self.env.seed(seed)
            self.env = SMACEnv(
                map_name=map_name
                )
    
    def reset(self):
        with silence_stderr():
            local_obs, global_state = self.env.reset()
            available_actions = []
            for i in range(self.env.n_agents):
                available_actions.append(self.env.get_avail_agent_actions(i))
            return local_obs, global_state, available_actions

    def step(self, actions):
        with silence_stderr():
            reward, done, info = self.env.step(actions)
            available_actions = []
            rewards = []
            dones = []
            infos = []
            info["won"] = self.env.win_counted
            for i in range(self.env.n_agents):
                available_actions.append(self.env.get_avail_agent_actions(i))
                rewards.append(reward)
                dones.append(done)
                infos.append(info)
            return self.env.get_obs(), self.env.get_state(), rewards, dones, infos, available_actions
    
    def close(self):
        with silence_stderr():
            self.env.close()